from zhipuai import ZhipuAI

class WebSearch:

    def __init__(self):
        self._client = ZhipuAI(api_key="51109e5e1fd67c96a7a51eb74e5ae8ca.SCpBvbMmVT6i3axx")

        self._tools = [{
            "type": "web_search",
            "web_search": {
                "enable": True #默认为关闭状态（False） 禁用：False，启用：True。
            }
        }]
    
    def __call__(self,query):
        messages = [{
            "role": "user",
            "content": query
        }]

        response = self._client.chat.completions.create(
            model="glm-4",
            messages=messages,
            tools=self._tools
        )
        return response.choices[0].message.content

websearch = WebSearch()

from doc import _zhaoyuxuan,_sunyuexing
from langchain_core.documents import Document
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings

class RAG:

    def __init__(self):

        _emb = OllamaEmbeddings(
            base_url="http://192.168.10.13:60001",
            model="bge-m3"
        )
        
        _docs = [
            Document(page_content=_zhaoyuxuan),
            Document(page_content=_sunyuexing)
        ]

        
        _text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
        _splits = _text_splitter.split_documents(_docs)
        self._vectorstore = Chroma.from_documents(_splits,_emb)

    def __call__(self,query):
        return self._vectorstore.similarity_search_with_score(query=query,k=2)
    
retriever = RAG()